import os
import tempfile
import subprocess
import logging
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
import re
from docx import Document
from docx.shared import Pt
from docx.oxml.ns import qn

app = FastAPI()

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建临时文件夹
temp_dir = os.path.join(os.getcwd(), "temp")
os.makedirs(temp_dir, exist_ok=True)

# LaTeX模板文件，用于渲染数学公式
latex_template = r"""
\documentclass{article}
\usepackage[utf8]{inputenc}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{unicode-math}
\usepackage{siunitx}
\usepackage[version=4]{mhchem}
\usepackage{chemformula}
\usepackage{upgreek}
\usepackage{textcomp}
\usepackage{gensymb}
\sisetup{
  per-mode = symbol,
  separate-uncertainty = true,
  range-phrase = \text{~to~},
  range-units = single
}
\begin{document}
$body$
\end{document}
"""


template_path = os.path.join(temp_dir, "latex_template.tex")
with open(template_path, "w") as f:
    f.write(latex_template)


def create_reference_docx(path):
    document = Document()
    # 修改默认样式
    style = document.styles['Normal']
    font = style.font
    font.name = 'DejaVu Sans'
    font.size = Pt(11)
    
    # 修改默认样式以包含数学字体设置
    style.element.rPr.rFonts.set(qn('w:asciiTheme'), 'Cambria Math')
    style.element.rPr.rFonts.set(qn('w:eastAsiaTheme'), 'Cambria Math')
    style.element.rPr.rFonts.set(qn('w:hAnsiTheme'), 'Cambria Math')
    style.element.rPr.rFonts.set(qn('w:cstheme'), 'Cambria Math')
    
    document.save(path)
# 在主代码中调用此函数
reference_doc_path = os.path.join(os.getcwd(), "reference.docx")
create_reference_docx(reference_doc_path)

def latex_to_unicode(match):
    latex = match.group(1)

    
    # 专门处理化学式
    if r'\si{' in latex and re.search(r'[A-Z][a-z]?(_\{?\d+\}?)+', latex):
        # 提取所有元素和下标
        elements = re.findall(r'([A-Z][a-z]?)(_\{?(\d+)\}?)?', latex)
        result = ''
        for element, _, subscript in elements:
            result += element + (subscript if subscript else '')
        return result
    
    # 处理化学式
    if r'\mathrm{' in latex and '_' in latex:
        return re.sub(r'\\mathrm{([A-Z][a-z]?)_?{?(\d+)?}?}',
                      lambda m: f'{m.group(1)}_{m.group(2)}' if m.group(2) else f'{m.group(1)}',
                      latex)
    
    # 处理单位和数值
    elif r'\si{' in latex or r'\mathrm{' in latex:
        # 保留原始格式
        return latex
    
    # 处理上标
    elif '^' in latex:
        return latex
    
    # 处理下标
    elif '_' in latex:
        return latex
    
    # 处理特殊符号
    elif r'\underline' in latex:
        if r'\ge' in latex or r'\geqslant' in latex:
            return '≥' + latex.split('}')[-1].strip()
        elif r'\le' in latex or r'\leqslant' in latex:
            return '≤' + latex.split('}')[-1].strip()
    elif r'\ge' in latex or r'\geqslant' in latex:
        match = re.search(r'(?:\\ge|\\geqslant)\s*(\d+(?:\.\d+)?)', latex)
        if match:
            return f'≥{match.group(1)}'
        return '≥'
    elif r'\le' in latex or r'\leqslant' in latex:
        match = re.search(r'(?:\\le|\\leqslant)\s*(\d+(?:\.\d+)?)', latex)
        if match:
            return f'≤{match.group(1)}'
        return '≤'
    elif r'\infty' in latex:
        return '∞'
    elif '%' in latex:
        return latex
    return latex



def preprocess_markdown(content):
    # 处理化学式，包括下标部分
    content = re.sub(r'\\mathrm{([A-Z][a-z]?)_?{?(\d+)?}?}',
                     lambda m: f'\\mathrm{{{m.group(1)}}}_{{{m.group(2)}}}' if m.group(2) else f'\\mathrm{{{m.group(1)}}}', content)
    # 处理 IC_50
    content = re.sub(r'\$\\mathrm{(IC)_{([0-9]+)}}\$',
                 r'IC<sub>\2</sub>',
                 content)
    # 处理单位，保留 mg/kg 等复合单位
    content = re.sub(
        r'\\mathrm{([a-zA-Z]+(/[a-zA-Z]+)?)}', r'\\si{\1}', content)
    # 处理数值和单位之间的空格，包括复合单位
    content = re.sub(r'(\d+(?:\.\d+)?)\s*[~,:]\s*(\\mathrm|\\si){([a-zA-Z]+(/[a-zA-Z]+)?)}',
                     r'\1~\\si{\3}', content)
    # 处理特殊情况，如 150 mg/
    content = re.sub(r'(\d+(?:\.\d+)?)\s*[~,:]\s*(\\mathrm|\\si){([a-zA-Z]+/)}',
                     r'\1~\\si{\3}', content)
    # 处理范围
    content = re.sub(r'(\d+)\s*~\s*-\s*~\s*(\d+)', r'\1 -- \2', content)
    # 处理大于等于和小于等于符号
    content = content.replace(r'{\ge}', r'\ge')
    content = content.replace(r'{\geq}', r'\ge')
    content = content.replace(r'{\le}', r'\le')
    content = content.replace(r'{\leq}', r'\le')
    # 处理特殊操作符
    content = content.replace(r'\operatorname*{max}', r'\max')
    # 处理 AUC_{0-t}
    content = re.sub(
        r'\\mathrm{([A-Z]+)_{([0-9]-[a-z])}}', r'\\mathrm{\1}_{\2}', content)
    # 处理下划线和大于等于符号的组合
    content = re.sub(r'\$\\underline\{\\ge\}\s*(\d+)\$',
                     r'$\\geqslant \1$', content)
    content = re.sub(r'\$\\underline\{\\geqslant\}\s*(\d+)\$',
                     r'$\\geqslant \1$', content)
    
    # 处理大于等于符号
    content = re.sub(r'\$\{?\\ge(?:qslant)?\}?\s*(\d+(?:\.\d+)?)\$$',
                     r'$\\geqslant \1$', content)
    # 处理百分比
    content = re.sub(r'\$(\d+(?:\.\d+)?)\\?%\$', r'$\1\\%$', content)
    # 处理平均值和标准差
    content = re.sub(r'(\d+(?:\.\d+)?)\s*±\s*(\d+(?:\.\d+)?)',
                     r'\1 \\pm \2', content)
    # 处理 $150\,\mathrm{mg/}$ 情况
    content = re.sub(r'\$(\d+(?:\.\d+)?)\s*\\,?\\mathrm{([a-zA-Z]+/?)}\$',
                     r'$\1~\\si{\2}$', content)
    # 处理大于等于符号的显示问题
    content = re.sub(r'\$\{?\\ge(?:qslant)?\}?\s*(\d+(?:\.\d+)?)\s*\\mathrm\{([a-zA-Z]+/[a-zA-Z]+)\}/?\$',
                     r'$\\geqslant \1~\\si{\2}$', content)
    # 最后将所有 LaTeX 数学符号转成 Unicode
    content = re.sub(r'\$(.*?)\$', lambda m: f'${latex_to_unicode(m)}$', content)
    return content

@app.post("/convert/")
async def convert_markdown_to_docx(input_file: UploadFile = File(...),
                                   output_file: str = "output.docx"):
    temp_markdown_fd, temp_markdown_path = tempfile.mkstemp(
        suffix=".md", dir=temp_dir)
    os.close(temp_markdown_fd)

    content = await input_file.read()
    logger.info(f"Original content: {content.decode()}")

    preprocessed_content = preprocess_markdown(content.decode())
    logger.info(f"Preprocessed content: {preprocessed_content}")

    with open(temp_markdown_path, "w", encoding="utf-8") as buffer:
        buffer.write(preprocessed_content)

    docx_path = os.path.join(temp_dir, output_file)
    if not os.path.exists(reference_doc_path):
        create_reference_docx(reference_doc_path)
    try:
        result = subprocess.run([
            "pandoc",
            temp_markdown_path,
            "-o", docx_path,
            "--from=markdown+tex_math_dollars+tex_math_single_backslash+latex_macros",
            "--to=docx",
            "--mathjax",
            "--standalone",
            "--wrap=none",
            "--citeproc", 
            "--pdf-engine=xelatex",
            f"--reference-doc={reference_doc_path}"
        ], check=True, capture_output=True, text=True)
        logger.info(f"Conversion successful. Output: {result.stdout}")
    except subprocess.CalledProcessError as e:
        logger.error(f"Conversion failed. Error: {e.stderr}")
        raise HTTPException(status_code=500, detail=f"Conversion failed: {e.stderr}")

    os.remove(temp_markdown_path)

    return FileResponse(docx_path, media_type='application/vnd.openxmlformats-officedocument.wordprocessingml.document', filename=output_file)
